//
// Created by zhangcc on 2021/5/24.
//
#include <math.h>
#include <float.h>

#include "trilinear_kernel.h"

int TriLinearForwardCpu(const float* lut, const float* src, float* dst, const int dim, \
                        const int shift, const float binsize, const int width, const int height, \
                        const int channels, const int bit_num)
{
    int shift_ = (shift << 1);
    int step = width * channels;

    const float scale = (1.0f / (float) ( (1 << bit_num) - 1) ) * (MAX_LEVEL - 1);

    const float* src_ptr = src;
    float* dst_ptr = dst;

    for (int row = 0; row < height; row++)
    {
        for (int col = 0; col < width * step; col += step)
        {
            float r = src_ptr[0] * scale;
            float g = src_ptr[1] * scale;
            float b = src_ptr[2] * scale;

            int r_id = int(r * (1 << MAX_SHIFT_NUM)) >> MAX_SHIFT_NUM;
            int g_id = int(g * (1 << MAX_SHIFT_NUM)) >> MAX_SHIFT_NUM;
            int b_id = int(b * (1 << MAX_SHIFT_NUM)) >> MAX_SHIFT_NUM;

            float r_d = (r - r_id);
            float g_d = (g - g_id);
            float b_d = (b - b_id);

            int id000 = r_id + g_id * dim + b_id * dim * dim;
            int id100 = r_id + 1 + g_id * dim + b_id * dim * dim;
            int id010 = r_id + (g_id + 1) * dim + b_id * dim * dim;
            int id110 = r_id + 1 + (g_id + 1) * dim + b_id * dim * dim;
            int id001 = r_id + g_id * dim + (b_id + 1) * dim * dim;
            int id101 = r_id + 1 + g_id * dim + (b_id + 1) * dim * dim;
            int id011 = r_id + (g_id + 1) * dim + (b_id + 1) * dim * dim;
            int id111 = r_id + 1 + (g_id + 1) * dim + (b_id + 1) * dim * dim;

            float w000 = (1 - r_d) * (1 - g_d) * (1 - b_d);
            float w100 = r_d * (1 - g_d) * (1 - b_d);
            float w010 = (1 - r_d) * g_d * (1 - b_d);
            float w110 = r_d * g_d * (1 - b_d);
            float w001 = (1 - r_d) * (1 - g_d) * b_d;
            float w101 = r_d * (1 - g_d) * b_d;
            float w011 = (1 - r_d) * g_d * b_d;
            float w111 = r_d * g_d * b_d;

            float out = w000 * lut[id000] + w100 * lut[id100] +
                        w010 * lut[id010] + w110 * lut[id110] +
                        w001 * lut[id001] + w101 * lut[id101] +
                        w011 * lut[id011] + w111 * lut[id111];
            dst_ptr[0] = out * ( (1 << bit_num) - 1);

            out = w000 * lut[id000 + shift] + w100 * lut[id100 + shift] +
                  w010 * lut[id010 + shift] + w110 * lut[id110 + shift] +
                  w001 * lut[id001 + shift] + w101 * lut[id101 + shift] +
                  w011 * lut[id011 + shift] + w111 * lut[id111 + shift];
            dst_ptr[1] = out * ( (1 << bit_num) - 1);

            out = w000 * lut[id000 + shift_] + w100 * lut[id100 + shift_] +
                  w010 * lut[id010 + shift_] + w110 * lut[id110 + shift_] +
                  w001 * lut[id001 + shift_] + w101 * lut[id101 + shift_] +
                  w011 * lut[id011 + shift_] + w111 * lut[id111 + shift_];
            dst_ptr[2] = out * ( (1 << bit_num) - 1);

            src_ptr += channels;
            dst_ptr += channels;
        }
    }

    return 0;
}